[ckpt] feat: implement large tensor slicing in vllm rollout and CheckpointEngine for weight updating by jianjunzhong · Pull Request #5378 · verl-project/verl

jianjunzhong · 2026-02-24T03:22:14Z

What does this PR do?

This PR does the following:

Adds support for large tensors slicing in the naive(vllm)/nccl/hccl/nixl backends of CheckpointEngine.
Abstracts common logic from nccl/hccl/nixl into the CheckpointEngine and CollectiveCheckpointEngine classes.
Adds a unit test for the correctness of naive weight updates.

Checklist Before Starting

Search for similar PRs. Paste at least one query link here: ...
Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
- {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward
- If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
- {type} is in feat, fix, refactor, chore, test
- If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
- Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

Read the Contribute Guide.
Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
Add / Update the documentation.
Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

…rge tensors Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

gemini-code-assist

Code Review

This pull request introduces chunked weight handling for large tensors in vLLM rollouts, which is a valuable feature for working with large models. The changes involve slicing oversized weights on the sender side and reassembling them on the receiver side. A new test is also added to verify the correctness of this new functionality. My review identified a critical bug in the chunking logic that could lead to a ZeroDivisionError, and a high-severity issue in the new test related to a hardcoded, user-specific file path. I have provided code suggestions to address both of these points. Apart from these issues, the implementation appears solid.

verl/workers/rollout/vllm_rollout/vllm_rollout.py

gemini-code-assist · 2026-02-24T03:25:03Z

tests/checkpoint_engine/test_naive_correctness.py

+    config.trainer.n_gpus_per_node = 8
+    config.trainer.nnodes = 1
+    config.actor_rollout_ref.actor.use_dynamic_bsz = True
+    config.actor_rollout_ref.model.path = os.path.expanduser("~/models/Qwen/Qwen3-VL-2B-Instruct")


The model path is hardcoded to a user-specific local directory. This makes the test non-portable and difficult for other developers to run. It's recommended to use a model from the Hugging Face Hub that can be downloaded automatically, or at least make this path configurable via an environment variable. Using a smaller model for this test would also make it faster and more efficient.

Suggested change

config.actor_rollout_ref.model.path = os.path.expanduser("~/models/Qwen/Qwen3-VL-2B-Instruct")

config.actor_rollout_ref.model.path = os.path.expanduser(os.environ.get("VERL_TEST_MODEL_PATH", "hf-internal-testing/tiny-random-LlamaForCausalLM"))

…isionError Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

wuxibin89 · 2026-02-24T06:06:08Z

CheckpointEngine should also slice large tensor as well.

jianjunzhong · 2026-02-24T06:13:14Z

CheckpointEngine should also slice large tensor as well.

OK, I will also implement large tensor slicing in CheckpointEngine.

pengwu22 · 2026-02-24T20:34:02Z

should merge after #5309
so that the newly-added logic can be thoroughly unit-tested.

…rge tensors Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…isionError Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…ort large tensors Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…ode reuse Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…hunked weight handling Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…class method Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…CL checkpoint engines Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…ne configs and clone tensors in pending chunks Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…t engines Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

[ckpt] feat: implement large tensor slicing in CheckpointEngine

jianjunzhong · 2026-02-25T06:52:37Z

should merge after #5309 so that the newly-added logic can be thoroughly unit-tested.

Thank you for the suggestion. Since similar logic has already been implemented and tested in CheckpointEngine, I think it would be reasonable to merge this PR first.

gemini-code-assist

Code Review

This pull request introduces support for large tensor slicing during weight updates in the CheckpointEngine and vllm rollout, refactors common logic into a new CollectiveCheckpointEngine base class, and adds a correctness test for the naive backend. While these changes enhance functionality and code structure, a high-severity security issue was identified: the use of pickle (via zmq.recv_pyobj) for deserializing metadata over the network poses a significant risk of Remote Code Execution (RCE). It is strongly recommended to switch to a safer serialization format like JSON. Additionally, there is a significant piece of duplicated code for the tensor slicing logic that could benefit from refactoring to improve long-term maintainability.

gemini-code-assist · 2026-02-25T06:55:46Z

verl/checkpoint_engine/base.py

+            self.socket.send_pyobj(self.metadata)
+        else:
+            self.socket.recv_string()
+            self.metadata = self.socket.recv_pyobj()


The use of zmq.Socket.recv_pyobj() is insecure as it relies on the pickle module for deserialization. pickle is known to be vulnerable to arbitrary code execution if the input data is controlled by an attacker. Since this data is received over the network from the master process, an attacker who can connect to the ZMQ port or spoof the master can achieve Remote Code Execution (RCE) on the worker processes.

Recommendation: Use a safer serialization format such as JSON or msgpack. For non-serializable types like torch.Size or torch.dtype, convert them to strings or lists before sending and reconstruct them on the receiving end.

verl/workers/rollout/vllm_rollout/vllm_rollout.py

…cing in ServerAdapter and CheckpointEngine Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

jianjunzhong · 2026-02-25T08:32:01Z

I think this PR is ready for review. cc @wuxibin89 @pengwu22

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

…ass configuration Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

jianjunzhong requested review from PeterSH6, chenhaiq and wuxibin89 as code owners February 24, 2026 03:22

[vllm] feat: implement chunked weight handling in vllm rollout for la…

fe7e78e

…rge tensors Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

jianjunzhong force-pushed the feat/chunked_weight_update branch from 5d76fbc to fe7e78e Compare February 24, 2026 03:24

gemini-code-assist bot reviewed Feb 24, 2026

View reviewed changes

fix: handle large weight tensors in chunking logic to prevent ZeroDiv…

46345a8

…isionError Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

jianjunzhong and others added 11 commits February 25, 2026 13:54

[vllm] feat: implement chunked weight handling in vllm rollout for la…

fcf6ecd

…rge tensors Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

fix: handle large weight tensors in chunking logic to prevent ZeroDiv…

9947884

…isionError Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

feat: implement chunked weight handling in checkpoint engines to supp…

9e1be9a

…ort large tensors Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

refactor: move tensor yielding logic to CheckpointEngine for better c…

b13a546

…ode reuse Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

feat: implement collective checkpoint engine for NCCL and HCCL with c…

9d0a80e

…hunked weight handling Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

refactor: implement chunk slicing in NIXLCheckpointEngine using base …

418fbd0

…class method Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

fix: add properties for bucket size, rank, and buffers in HCCL and NC…

ca5c4aa

…CL checkpoint engines Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

fix: add update_weights_bucket_megabytes parameter to checkpoint engi…

2a84ae4

…ne configs and clone tensors in pending chunks Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

refactor: streamline HCCL and NCCL checkpoint engines

4d617da

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

fix: update total_params calculation in NIXL and Collective checkpoin…

177c9f6

…t engines Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

[ckpt] feat: implement large tensor slicing in CheckpointEngine

d781221

[ckpt] feat: implement large tensor slicing in CheckpointEngine

jianjunzhong changed the title ~~[vllm] feat: implement chunked weight handling in vllm rollout for large tensors~~ [ckpt] feat: implement large tensor slicing in vllm rollout and CheckpointEngine for weight updating Feb 25, 2026

gemini-code-assist bot reviewed Feb 25, 2026

View reviewed changes

jianjunzhong added 2 commits February 25, 2026 15:44

refactor: utilize compute_weight_chunks utility for weight tensor sli…

ec18f2b

…cing in ServerAdapter and CheckpointEngine Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

fix: add bucket_size property with getter and setter

ec908ea

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

jianjunzhong added 3 commits February 26, 2026 13:09

Merge branch 'main' into feat/chunked_weight_update

5af5eed

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

refactor: update CheckpointEngineManager initialization to use datacl…

d31bb52

…ass configuration Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

refactor: move logger initialization to the top of tensor_utils.py

80e3235

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

test: temporarily skip test_naive_correctness.py

e7ed83a

Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ckpt] feat: implement large tensor slicing in vllm rollout and CheckpointEngine for weight updating#5378

[ckpt] feat: implement large tensor slicing in vllm rollout and CheckpointEngine for weight updating#5378
jianjunzhong wants to merge 19 commits intoverl-project:mainfrom
jianjunzhong:feat/chunked_weight_update

jianjunzhong commented Feb 24, 2026 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

gemini-code-assist bot Feb 24, 2026

Uh oh!

wuxibin89 commented Feb 24, 2026

Uh oh!

jianjunzhong commented Feb 24, 2026

Uh oh!

pengwu22 commented Feb 24, 2026

Uh oh!

jianjunzhong commented Feb 25, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Feb 25, 2026

Uh oh!

Uh oh!

jianjunzhong commented Feb 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	config.actor_rollout_ref.model.path = os.path.expanduser("~/models/Qwen/Qwen3-VL-2B-Instruct")
	config.actor_rollout_ref.model.path = os.path.expanduser(os.environ.get("VERL_TEST_MODEL_PATH", "hf-internal-testing/tiny-random-LlamaForCausalLM"))

Conversation

jianjunzhong commented Feb 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Checklist Before Starting

Test

API and Usage Example

Design & Code Changes

Checklist Before Submitting

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist bot Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

wuxibin89 commented Feb 24, 2026

Uh oh!

jianjunzhong commented Feb 24, 2026

Uh oh!

pengwu22 commented Feb 24, 2026

Uh oh!

jianjunzhong commented Feb 25, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Feb 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jianjunzhong commented Feb 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jianjunzhong commented Feb 24, 2026 •

edited

Loading